import numpy as np
import cv2
import math
import imageio
import os
import datetime
import pandas as pd
from lib import Heading

# 配置参数
class Config:
    def __init__(self):
        self.MAX_VEL = 1.5  # 最大线速度
        self.MIN_VEL = 0.0  # 最小线速度
        self.MAX_ANG_VEL = np.pi / 2  # 最大角速度
        self.MAX_ACCEL = 1.5  # 最大线加速度
        self.MAX_ANG_ACCEL = np.pi / 1  # 最大角加速度
        self.DT = 0.1  # 时间间隔
        self.V_RESOLUTION = 0.02  # 线速度分辨率
        self.ANG_V_RESOLUTION = np.pi / 180  # 角速度分辨率
        self.PREDICT_TIME = 1.5  # 预测时间
        self.GOAL_DISTANCE = 0.5  # 到达目标的距离阈值
        self.OBSTACLE_DISTANCE = 0.5  # 障碍物安全距离
        self.GOAL_COST_GAIN = 0.3  # 目标距离权重
        self.SPEED_COST_GAIN = 0.3  # 速度权重
        self.OBSTACLE_COST_GAIN = 0.4  # 障碍物距离权重
        self.DYNAMIC_COST_GAIN = 0.01
        self.OBSTACLE_PREDICT_TIME = 0
        self.robot_radius = 0.1
        self.OBSTACLE_RADIUS = 0.1
        self.scale = 50
        self.canvas_size = (600, 600)


config = Config()
# 创建输出文件夹
now_str = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
out_dir = f"DWA_improve_Output_{now_str}"
os.makedirs(out_dir, exist_ok=True)

# 初始状态和目标
x = np.array([0.0, 0.0, np.pi/4, 1.5, 0.0])  # [x, y, theta, v, omega]
goal = np.array([10.0, 10.0])

# 障碍物 [x, y, vx, vy]
obstacles = np.array([
    [4.5, 5, 0.1, 0],
    [6, 8, -0.05, -0.05],
    [7, 5, 0, 0.1]
])

trajectory = [x.copy()]
candidate_trajs = []
best_traj_global = None
frames = []  # 存储GIF帧


def motion(x, u, dt):
    """根据当前状态和控制输入更新机器人状态"""
    x_new = x.copy()
    x_new[2] += u[1] * dt
    x_new[0] += u[0] * np.cos(x_new[2]) * dt
    x_new[1] += u[0] * np.sin(x_new[2]) * dt
    x_new[3] = u[0]
    x_new[4] = u[1]
    return x_new


def calc_dynamic_window(x):
    """计算动态窗口 """
    # 机器人运动学约束
    vs = [config.MIN_VEL, config.MAX_VEL, -config.MAX_ANG_VEL, config.MAX_ANG_VEL]

    # 加速度约束
    vd = [x[3] - config.MAX_ACCEL * config.DT, x[3] + config.MAX_ACCEL * config.DT,
          x[4] - config.MAX_ANG_ACCEL * config.DT, x[4] + config.MAX_ANG_ACCEL * config.DT]

    # 动态窗口
    dw = [max(vs[0], vd[0]), min(vs[1], vd[1]),
          max(vs[2], vd[2]), min(vs[3], vd[3])]

    return dw


def predict_trajectory(x_init, v, omega):
    """预测机器人在给定速度和角速度下的轨迹"""
    x = np.array(x_init)
    trajectory = np.array(x)
    time = 0
    while time <= config.PREDICT_TIME:
        x = motion(x, [v, omega], config.DT)
        trajectory = np.vstack((trajectory, x))
        time += config.DT
    return trajectory


def update_obstacles(obstacles, dt):
    """更新障碍物的位置 """
    for i in range(len(obstacles)):
        obstacles[i][0] += obstacles[i][2] * dt
        obstacles[i][1] += obstacles[i][3] * dt
    return obstacles

def predict_obstacle_trajectory(obstacles, predict_time):
    """预测障碍物的位置 """
    future_obstacles = []
    for obs in obstacles:
        obs_x, obs_y, vx, vy = obs
        future_x = obs_x + vx * predict_time
        future_y = obs_y + vy * predict_time
        future_obstacles.append([future_x, future_y])
    return np.array(future_obstacles)

def calc_cost_goal(trajectory, goal):
    dx = goal[0] - trajectory[-1, 0]
    dy = goal[1] - trajectory[-1, 1]
    goal_dist = np.sqrt(dx ** 2 + dy ** 2)
    return config.GOAL_COST_GAIN * goal_dist

def calc_cost_obstacle(trajectory, ob):
    ob_dist = float("inf")
    for i in range(len(trajectory)):
        for j in range(len(ob)):
            dx = trajectory[i, 0] - ob[j, 0]
            dy = trajectory[i, 1] - ob[j, 1]
            d = np.sqrt(dx ** 2 + dy ** 2)
            if d < ob_dist:
                ob_dist = d
    if ob_dist <= config.OBSTACLE_DISTANCE:
        obstacle_cost = config.OBSTACLE_COST_GAIN * (1 / ob_dist)
    else:
        obstacle_cost = 0
    return obstacle_cost

def calc_cost_speed(trajectory):
    return config.SPEED_COST_GAIN * (config.MAX_VEL - trajectory[-1, 3])

def calc_dynamic_obstacle_cost(trajectory, obstacles):
    cost = 0.0
    for obs in obstacles:
        obs_x, obs_y, vx, vy = obs
        obs_heading = np.arctan2(vy, vx)  # 障碍物的运动方向
        dx = trajectory[-1, 0] - obs_x
        dy = trajectory[-1, 1] - obs_y
        DIS = np.sqrt(dx ** 2 + dy ** 2)
        if DIS < 2.5:
            robo_heading = trajectory[-1, 2]
            angle_diff = robo_heading - obs_heading
            angle_diff = (angle_diff + np.pi) % (2 * np.pi) - np.pi
            if np.abs(np.cos(angle_diff)) > 0.5:
                cost += config.DYNAMIC_COST_GAIN
    return cost

def dwa_control(x, goal, obstacles):
    dw = calc_dynamic_window(x)
    best_u = [0.0, 0.0]
    min_cost = float("inf")
    best_trajectory = np.array(x)
    all_trajectories = []

    for v in np.arange(dw[0], dw[1], config.V_RESOLUTION):
        for omega in np.arange(dw[2], dw[3], config.ANG_V_RESOLUTION):
            trajectory = predict_trajectory(x, v, omega)
            future_obstacles = predict_obstacle_trajectory(obstacles, config.OBSTACLE_PREDICT_TIME)
            cost = (calc_cost_goal(trajectory, goal) +
                    calc_cost_obstacle(trajectory, future_obstacles) +
                    calc_cost_speed(trajectory) +
                    calc_dynamic_obstacle_cost(trajectory, obstacles))
            all_trajectories.append(trajectory)
            if cost < min_cost:
                min_cost = cost
                best_u = [v, omega]
                best_trajectory = trajectory

    return best_u, best_trajectory, all_trajectories


def draw(canvas, robot, goal, obstacles, traj, candidate_trajs, best_traj):
    canvas[:] = 255
    for obs in obstacles:
        obs_px = (int(obs[0] * config.scale), int(obs[1] * config.scale))
        r = int((config.OBSTACLE_RADIUS) * config.scale)
        obs_safe = int((config.OBSTACLE_DISTANCE) * config.scale)
        cv2.circle(canvas, obs_px, r, (0, 0, 0), -1)
        cv2.circle(canvas, obs_px, obs_safe, (0, 0, 0), 2)  # Safe distance

    for cand in candidate_trajs:
        for i in range(1, len(cand)):
            p1 = (int(cand[i - 1][0] * config.scale), int(cand[i - 1][1] * config.scale))
            p2 = (int(cand[i][0] * config.scale), int(cand[i][1] * config.scale))
            cv2.line(canvas, p1, p2, (200, 200, 200), 1)

    if best_traj is not None:
        for i in range(1, len(best_traj)):
            p1 = (int(best_traj[i - 1][0] * config.scale), int(best_traj[i - 1][1] * config.scale))
            p2 = (int(best_traj[i][0] * config.scale), int(best_traj[i][1] * config.scale))
            cv2.line(canvas, p1, p2, (0, 0, 255), 2)

    # Draw actual trajectory
    for i in range(1, len(traj)):
        p1 = (int(traj[i - 1][0] * config.scale), int(traj[i - 1][1] * config.scale))
        p2 = (int(traj[i][0] * config.scale), int(traj[i][1] * config.scale))
        cv2.line(canvas, p1, p2, (255, 0, 0), 2)

    # Draw goal and robot
    cv2.circle(canvas, (int(goal[0] * config.scale), int(goal[1] * config.scale)), 5, (0, 0, 255), -1)
    robot_px = (int(robot[0] * config.scale), int(robot[1] * config.scale))
    cv2.circle(canvas, robot_px, int(config.robot_radius * config.scale), (255, 0, 255), 2)


# 主循环 (from DWA_improve with DWA_2 logic)
canvas = np.ones((config.canvas_size[1], config.canvas_size[0], 3), dtype=np.uint8) * 255
heading_records = []
speed_records = []
coordinate_x = []
coordinate_y = []
path_length = 0.0  # 初始化路径长度

while np.linalg.norm(x[:2] - goal) > config.GOAL_DISTANCE:
    obstacles = update_obstacles(obstacles, config.DT)
    ob = obstacles[:, :2]
    prev_position = x[:2].copy()
    u, best_traj, candidate_trajs = dwa_control(x, goal, obstacles)
    x = motion(x, u, config.DT)
    trajectory.append(x.copy())
    current_position = x[:2]
    path_length += np.linalg.norm(current_position - prev_position)
    # Record data for analysis
    heading_records.append(x[2] * 57.3)  # Convert to degrees
    speed_records.append(x[3])
    coordinate_x.append(x[0])
    coordinate_y.append(x[1])
    # Visualization
    draw(canvas, x, goal, obstacles, trajectory, candidate_trajs, best_traj)
    frame = cv2.cvtColor(canvas.copy(), cv2.COLOR_BGR2RGB)
    frames.append(frame)
    cv2.imshow('DWA_improve Simulation', canvas)
    key = cv2.waitKey(1)
    if key == 27:  # ESC key to break
        break

# 导出GIF
gif_path = os.path.join(out_dir, "trajectory.gif")
imageio.mimsave(gif_path, frames, duration=0.05)

# 分析路径平滑度
evaluator = Heading.PathSmoothnessEvaluator(heading_records, unit='deg')
result = evaluator.evaluate()
print(f"路径总长度: {path_length}")
print(f"累计转角: {result['total_turning_angle']:.2f}°")
print(f"最大单次转角: {result['max_turning_angle']:.2f}°")

# 导出Excel
timestamps = [i for i in range(len(heading_records))]
df_records = pd.DataFrame({
    'Time': timestamps,
    'X': coordinate_x,
    'Y': coordinate_y,
    'Speed': speed_records,
    'Yaw': heading_records
})
df_records.to_excel(os.path.join(out_dir, "velocity_heading.xlsx"), index=False)

cv2.destroyAllWindows()
print(f"✅ 动图和航速航向数据已保存至文件夹：{out_dir}")